import sys
import os
import requests
from PyQt5.QtCore import QUrl, Qt
from PyQt5.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout,
    QLineEdit, QPushButton, QToolBar, QMessageBox,
    QTabWidget, QAction
)
from PyQt5.QtWebEngineWidgets import QWebEngineView, QWebEngineProfile, QWebEngineDownloadItem, QWebEnginePage

class CustomWebPage(QWebEnginePage):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.browser_window = None
    
    def set_browser_window(self, window):
        self.browser_window = window
    
    def acceptNavigationRequest(self, url, navigation_type, is_main_frame):
        # 处理不同类型的导航请求
        # 使用数字常量表示导航类型：NavigationTypeLinkClicked = 1
        if navigation_type == 1:  # NavigationTypeLinkClicked
            # 对于链接点击，如果是主框架，允许在当前标签页导航
            if is_main_frame:
                return True
            # 如果不是主框架（可能是新开窗口请求），在新标签页打开
            else:
                if self.browser_window:
                    new_tab = self.browser_window.create_tab(url.toString())
                    return False  # 阻止默认行为，因为我们已经处理了
        
        # 其他类型的导航（如表单提交、历史记录导航等）允许默认行为
        return super().acceptNavigationRequest(url, navigation_type, is_main_frame)
    
    def createWindow(self, window_type):
        # 当需要创建新窗口时（如window.open()调用），创建新标签页
        if self.browser_window:
            new_tab = self.browser_window.create_tab()
            return new_tab.page()
        return super().createWindow(window_type)


class SimpleWebBrowser(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Simple Web Browser")
        self.setGeometry(100, 100, 1200, 700)
        
        # Create downloads directory
        self.download_dir = os.path.join(os.getcwd(), "downloaded_images")
        if not os.path.exists(self.download_dir):
            os.makedirs(self.download_dir)
        
        # Create the UI elements
        self.create_widgets()
        
        # Set up download handling
        self.setup_download_handling()
        
        # Create initial tab
        self.create_tab("https://www.baidu.com")
        
        # Set up keyboard shortcuts
        self.setup_shortcuts()
    
    def create_widgets(self):
        # Create the toolbar
        toolbar = QToolBar()
        self.addToolBar(toolbar)
        
        # New tab button
        new_tab_action = QAction("+", self)
        new_tab_action.triggered.connect(lambda: self.create_tab())
        toolbar.addAction(new_tab_action)
        
        # Back button
        self.back_button = QPushButton("←")
        self.back_button.clicked.connect(self.go_back)
        toolbar.addWidget(self.back_button)
        
        # Forward button
        self.forward_button = QPushButton("→")
        self.forward_button.clicked.connect(self.go_forward)
        toolbar.addWidget(self.forward_button)
        
        # Refresh button
        self.refresh_button = QPushButton("↻")
        self.refresh_button.clicked.connect(self.refresh_page)
        toolbar.addWidget(self.refresh_button)
        
        # URL entry
        self.url_entry = QLineEdit()
        self.url_entry.returnPressed.connect(self.load_url)
        toolbar.addWidget(self.url_entry)
        
        # Go button
        self.go_button = QPushButton("Go")
        self.go_button.clicked.connect(self.load_url)
        toolbar.addWidget(self.go_button)
        
        # Download Images button
        self.download_images_button = QPushButton("Download Images")
        self.download_images_button.clicked.connect(self.download_all_images)
        toolbar.addWidget(self.download_images_button)
        
        # Create central widget and layout
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        layout = QVBoxLayout(central_widget)
        
        # Create tab widget
        self.tab_widget = QTabWidget()
        self.tab_widget.setTabsClosable(True)
        self.tab_widget.tabCloseRequested.connect(self.close_tab)
        self.tab_widget.currentChanged.connect(self.change_active_tab)
        layout.addWidget(self.tab_widget)
    
    def setup_shortcuts(self):
        # Create actions for keyboard shortcuts
        new_tab_shortcut = QAction("New Tab", self)
        new_tab_shortcut.setShortcut("Ctrl+T")
        new_tab_shortcut.triggered.connect(lambda: self.create_tab())
        self.addAction(new_tab_shortcut)
        
        close_tab_shortcut = QAction("Close Tab", self)
        close_tab_shortcut.setShortcut("Ctrl+W")
        close_tab_shortcut.triggered.connect(lambda: self.close_current_tab())
        self.addAction(close_tab_shortcut)
    
    def load_url(self):
        url = self.url_entry.text().strip()
        
        # Check if URL is empty
        if not url:
            return
            
        # Add protocol if missing
        if not url.startswith(("http://", "https://")):
            url = "https://" + url
            
        # Load URL in current web view
        current_web_view = self.tab_widget.currentWidget()
        if current_web_view:
            current_web_view.load(QUrl(url))
    
    def update_url_bar(self, q):
        self.url_entry.setText(q.toString())
        self.url_entry.setCursorPosition(0)
    
    def create_tab(self, url="https://www.baidu.com"):
        # Create a new web view
        web_view = QWebEngineView()
        
        # Create custom web page and set it to the web view
        custom_page = CustomWebPage()
        custom_page.set_browser_window(self)
        web_view.setPage(custom_page)
        
        # Handle URL changes
        web_view.urlChanged.connect(self.update_url_bar)
        
        # Load the URL
        web_view.load(QUrl(url))
        
        # Add the web view to a new tab
        tab_index = self.tab_widget.addTab(web_view, "New Tab")
        self.tab_widget.setCurrentIndex(tab_index)
        
        # Update tab title when page title changes
        web_view.titleChanged.connect(
            lambda title, idx=tab_index: self.tab_widget.setTabText(idx, 
            title[:30] + '...' if title and len(title) > 30 else (title or "New Tab"))
        )
        
        return web_view
    

    
    def close_tab(self, index):
        # Don't close the last tab
        if self.tab_widget.count() > 1:
            self.tab_widget.removeTab(index)
    
    def close_current_tab(self):
        # Close the current tab
        current_index = self.tab_widget.currentIndex()
        self.close_tab(current_index)
    
    def change_active_tab(self, index):
        # Update the URL bar when changing tabs
        if index >= 0:
            web_view = self.tab_widget.widget(index)
            if web_view:
                self.update_url_bar(web_view.url())
    
    def go_back(self):
        current_web_view = self.tab_widget.currentWidget()
        if current_web_view:
            current_web_view.back()
    
    def go_forward(self):
        current_web_view = self.tab_widget.currentWidget()
        if current_web_view:
            current_web_view.forward()
    
    def refresh_page(self):
        current_web_view = self.tab_widget.currentWidget()
        if current_web_view:
            current_web_view.reload()
    
    def setup_download_handling(self):
        # Connect download signals
        profile = QWebEngineProfile.defaultProfile()
        profile.downloadRequested.connect(self.download_requested)
    
    def download_requested(self, download):
        # Handle download requests
        path = os.path.join(self.download_dir, download.suggestedFileName())
        download.setPath(path)
        download.accept()
        
        # Connect download finished signal
        download.finished.connect(lambda: self.download_finished(download))
    
    def download_finished(self, download):
        # In PyQt5, we can check download state directly
        if download.state() == 3:  # QWebEngineDownloadItem.DownloadCompleted
            QMessageBox.information(self, "Download Completed", 
                                  f"File downloaded to:\n{download.path()}")
        elif download.state() == 4:  # QWebEngineDownloadItem.DownloadInterrupted
            QMessageBox.warning(self, "Download Error", 
                              f"Failed to download file: {download.url().toString()}")
    
    def download_all_images(self):
        # Execute JavaScript to get all image URLs from the page
        script = """
        var images = document.getElementsByTagName('img');
        var imageUrls = [];
        for (var i = 0; i < images.length; i++) {
            var src = images[i].src;
            if (src && src.startsWith('http')) {
                imageUrls.push(src);
            }
        }
        imageUrls;
        """
        
        current_web_view = self.tab_widget.currentWidget()
        if current_web_view:
            current_web_view.page().runJavaScript(script, self.process_image_urls)
    
    def process_image_urls(self, image_urls):
        if not image_urls:
            QMessageBox.information(self, "No Images", "No images found on this page.")
            return
        
        # Create a directory for images if it doesn't exist
        if not os.path.exists(self.download_dir):
            os.makedirs(self.download_dir)
        
        # Download each image
        downloaded_count = 0
        for i, url in enumerate(image_urls):
            try:
                # Skip data URLs (base64 encoded images)
                if url.startswith('data:image'):
                    continue
                
                # Generate a filename
                filename = f"image_{i+1}.jpg"
                if '/' in url:
                    url_filename = url.split('/')[-1].split('?')[0]
                    if '.' in url_filename:
                        filename = url_filename
                
                filepath = os.path.join(self.download_dir, filename)
                
                # Download the image
                response = requests.get(url, timeout=10)
                if response.status_code == 200:
                    with open(filepath, 'wb') as f:
                        f.write(response.content)
                    downloaded_count += 1
            except Exception as e:
                print(f"Error downloading image {url}: {str(e)}")
                continue
        
        # Show completion message
        if downloaded_count > 0:
            QMessageBox.information(self, "Download Completed", 
                                  f"Successfully downloaded {downloaded_count} images to {self.download_dir}")


def main():
    app = QApplication(sys.argv)
    browser = SimpleWebBrowser()
    browser.show()
    sys.exit(app.exec_())

if __name__ == "__main__":
    main()
